package com.gxnu.edu.cn.ttxxt.webSocket;

import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class WebSocketHandler extends TextWebSocketHandler {
    private static final ConcurrentHashMap<String, WebSocketSession> onlineUsers = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, String> userIps = new ConcurrentHashMap<>();

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        String payload = message.getPayload();
        String[] parts = payload.split(":", 3); // 解析协议：TYPE:SENDER:CONTENT

        if (parts.length >= 3) {
            String type = parts[0];
            String sender = parts[1];
            String content = parts[2];

            switch (type) {
                case "MESSAGE":
                    broadcastMessage(String.format("MESSAGE:%s:%s", sender, content), sender);
                    break;
                case "SYSTEM":
                    broadcastMessage(String.format("SYSTEM:%s", content), null);
                    break;
            }
        }
    }

    private void broadcastMessage(String message, String excludeUser) throws IOException {
        String senderIp = userIps.get(excludeUser);
        for (Map.Entry<String, WebSocketSession> entry : onlineUsers.entrySet()) {
            if (!entry.getKey().equals(excludeUser)) {
                String receiverIp = userIps.get(entry.getKey());
                if (senderIp != null && receiverIp != null && isSameIntranet(senderIp, receiverIp)) {
                    entry.getValue().sendMessage(new TextMessage(message));
                }
            }
        }
    }

    private boolean isSameIntranet(String ip1, String ip2) {
        String[] parts1 = ip1.split("\\.");
        String[] parts2 = ip2.split("\\.");
        if (parts1.length >= 3 && parts2.length >= 3) {
            return parts1[0].equals(parts2[0]) && parts1[1].equals(parts2[1]) && parts1[2].equals(parts2[2]);
        }
        return false;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String username = (String) session.getAttributes().get("username");
        String ip = (String) session.getAttributes().get("ip");
        if (username != null) {
            onlineUsers.put(username, session);
            userIps.put(username, ip);
            broadcastMessage(String.format("SYSTEM:%s 加入了聊天室", username), null);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String username = (String) session.getAttributes().get("username");
        if (username != null) {
            onlineUsers.remove(username);
            userIps.remove(username);
            broadcastMessage(String.format("SYSTEM:%s 离开了聊天室", username), null);
        }
    }
}